import re
import json


def load_dataset(test_file, sid_list):
    dataset = []
    # 按照题型load
    # 再按照分类加分类
    with open(test_file, "r") as f:
        data = json.load(f)
    for item in data:
        if "A" in item:
            each_data = load_choice(item)
        else:
            each_data = load_generation(item)
        dataset += each_data
    if sid_list != "all":
        print("评测分类：", sid_list.split(","))
        filtered_dataset = []
        for item in dataset:
            if item["sid"] in sid_list.split(","):
                filtered_dataset.append(item)
        dataset = filtered_dataset
    # import random
    # random.seed(2025)
    # random.shuffle(dataset)
    # dataset = dataset[:20]
    return dataset


def remove_option(input_str):
    pattern = r"^[A-Za-z]\.\s*"
    result = re.sub(pattern, "", input_str)
    return result


def load_choice(item):
    if "D" in item:
        item["D"] = remove_option(item["D"])
    if "C" in item:
        item["C"] = remove_option(item["C"])
    if "B" in item:
        item["B"] = remove_option(item["B"])
    if "A" in item:
        item["A"] = remove_option(item["A"])
    if "D" in item and item["D"].strip() != "":
        input_data = "{question}\nA. {A}\nB. {B}\nC. {C}\nD. {D}\n".format(**item)
        system_prompt = "你正在参加期货行业从业资格考试，请你从A、B、C、D这几个选项中选出一个或多个作为回答问题的最恰当的答案，你只输出代表选项的字符，不要解释。"
    elif "C" in item and item["C"].strip() != "":
        input_data = "{question}\nA. {A}\nB. {B}\nC. {C}\n".format(**item)
        system_prompt = "你正在参加期货行业从业资格考试，请你从A、B、C这几个选项中选出一个或多个作为回答问题的最恰当的答案，你只输出代表选项的字符，不要解释。"
    else:
        # 判断题
        input_data = "{question}\nA. {A}\nB. {B}\n".format(**item)
        system_prompt = "你正在参加期货行业从业资格考试，请你从A、B选项中选出一个作为回答问题的最恰当的答案，你只能输出一个字符，并且这个字符是A、B中一个。不要解释。"
    prompt_constructor = system_prompt + "\n" + "问题：" + input_data
    prompt_constructor = prompt_constructor.replace("\\n", "\n")
    answer = "{answer}".format(**item)
    id_split = item["id"].split("-")
    domain = id_split[0]
    subject = id_split[1]
    sid = domain + "_" + subject
    result = {}
    result["uid"] = item["uid"]
    result["sid"] = sid
    result["subject"] = subject
    result["domain"] = domain
    result["prompt"] = prompt_constructor
    result["answer"] = answer
    return result


# def load_compliance(data):
#     dataset = list()
#     for data_dict in data:
#         prompt_constructor = '{question}'.format(**data_dict)
#         answer = 'nan'
#         id_split = data_dict['id'].split("-")
#         domain = id_split[0]
#         subject = id_split[1]
#         sid = domain + "_" + subject
#         details = {}
#         details["uid"] = data_dict['uid']
#         details["sid"] = sid
#         details["subject"] = subject
#         details["domain"] = domain
#         details["prompt"] = prompt_constructor
#         details["answer"] = answer
#         dataset.append(details)
#     return dataset

# def load_recognization(data):
#     dataset = list()
#     for data_dict in data:
#         prompt_constructor = '{question}'.format(**data_dict)
#         answer = '{answer}'.format(**data_dict)
#         id = data_dict['id']
#         details = {}
#         details["id"] = int(id)
#         details["subject"] = subject
#         details["domain"] = domain
#         details["prompt"] = prompt_constructor
#         details["answer"] = answer
#         dataset.append(details)
#     return dataset


def load_generation(item):
    """
    多个得分点用 'criterium{}': 分开
    只跑一次推理，评测时，将模型结果依次和得分点比较，看是否包含
    """
    prompt_constructor = "{question}".format(**item)
    prompt_constructor = prompt_constructor.replace("\\n", "\n")
    answer = dict()
    idx = 1
    while True:
        if "criterium{}".format(str(idx)) in item:
            answer["criterium{}".format(str(idx))] = item[
                "criterium{}".format(str(idx))
            ]
            idx += 1
        else:
            break
    id_split = item["id"].split("-")
    domain = id_split[0]
    subject = id_split[1]
    sid = domain + "_" + subject
    result = {}
    result["uid"] = item["uid"]
    result["sid"] = sid
    result["subject"] = subject
    result["domain"] = domain
    result["prompt"] = prompt_constructor
    result["answer"] = answer
    return result


def split_standard(domain_name):
    targets = ["_dev", "_val", "_test"]
    for item in targets:
        if domain_name.endswith(item):
            domain_name = domain_name.split(item)[0]
    return domain_name
